#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <iostream>
#include <cstdlib>
#include "common.h"
#include <vector>
#include <time.h>
#include <sys/time.h>
#include "test_kernel_launch.cuh" 

#define USECPSEC 1000000ULL

unsigned long long dtime_usec(unsigned long long start){

  timeval tv;
  gettimeofday(&tv, 0);
  return ((tv.tv_sec*USECPSEC)+tv.tv_usec)-start;
}

// perform vector averaging over M vectors of length L,  followed by matrix-vector multiply
// repeat the above N times
// input vectors are stored as a set of N column-major matrices
// for each k in N: output[k] = matrix*input[k]
template <typename T>
void cpu_version1(T *input, T *output, T *matrix, int L, int M, int N){
#pragma omp parallel for
  for (int k = 0; k < N; k++){      // repeat the following, N times
    std::vector<T> v1(L);           // vector length of L
    for (int i = 0; i < M; i++)     // compute average vector over M input vectors
      for (int j = 0; j < L; j++)
        v1[j] += input[k*M*L+j*M+i];
    for (int j = 0; j < L; j++)
      v1[j] /= M;
    for (int i = 0; i < L; i++)     // matrix-vector multiply
      for (int j = 0; j < L; j++)
	output[i*N+k] += matrix[i*L+j]*v1[j];
  }
}

// This will output the proper CUDA error strings in the event
// that a CUDA host call returns an error
#define checkGPUErrors(val) __check((val), #val, __FILE__, __LINE__)
#define cudaCheckErrors(msg) \
  do { \
    cudaError_t __err = cudaGetLastError(); \
    if (__err != cudaSuccess) { \
        fprintf(stderr, "Fatal error: %s (%s at %s:%d)\n", \
            msg, cudaGetErrorString(__err), \
            __FILE__, __LINE__); \
        fprintf(stderr, "*** FAILED - ABORTING\n"); \
        exit(1); \
    } \
  } while (0)

void __check(cudaError_t result, char const *const func, const char *const file,
           int const line) {
  if (result) {
    fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", file, line,
            static_cast<unsigned int>(result), cudaGetErrorName(result), func);
    exit(EXIT_FAILURE);
  }
}

  __global__ void warm_up()
{

}
struct TestResult {
    std::string filename;
    float average_time;
};
template <typename T>
float test_case(std::vector<int> input_size, std::vector<int> output_size,std::string  filename) {
    CUDA_INIT();
    //预热
    for(int i = 0; i < 5; i++) {
        warm_up<<<1,1,0,stream>>>();
    }
    
    //准备kernel输入数据
    ft *d_input, *h_input, *d_output, *h_outputc, *h_outputg, *d_matrix, *h_matrix;
    int L = my_L; int M = my_M; int N = my_N;
    // host allocations
    h_input   = new ft[N*L*M];
    h_matrix  = new ft[L*L];
    h_outputg = new ft[N*L];
    h_outputc = new ft[N*L];
    // data initialization
    for (int i = 0; i < N*L*M; i++) h_input[i] = (i % 2) +1;//(rand()&1)+1;  // 1 or 2
    for (int i = 0; i < L*L; i++) h_matrix[i]  = ((i+1) % 2) +1;//(rand()&1)+1;  // 1 or 2
    // create result to test for correctness
    unsigned long long dt = dtime_usec(0);
    cpu_version1(h_input, h_outputc, h_matrix, L, M, N);
    dt = dtime_usec(dt);
    std::cout << "CPU execution time: " << dt/(float)USECPSEC << "s" << std::endl;
    // device allocations
    cudaMalloc(&d_input, N*L*M*sizeof(ft));
    cudaMalloc(&d_output,  N*L*sizeof(ft));
    cudaMalloc(&d_matrix,  L*L*sizeof(ft));
    cudaCheckErrors("cudaMalloc failure");
    // copy input data from host to device
    cudaMemcpy(d_input,  h_input,  N*L*M*sizeof(ft), cudaMemcpyHostToDevice);
    cudaMemcpy(d_matrix, h_matrix,   L*L*sizeof(ft), cudaMemcpyHostToDevice);
    cudaMemset(d_output, 0, N*L*sizeof(ft));
    cudaCheckErrors("cudaMemcpy/Memset failure");

    //执行算子，记录时间
    float total_time = 0.0;
    int test_count = 10;
    float average_time = 0.0;
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    for (int i = 0; i < test_count; i++)
    {
        cudaEventRecord(start, 0);
        test_kernel_launch(d_input, d_output, d_matrix, L, M, N, stream);
        cudaEventRecord(stop, 0);
        cudaEventSynchronize(stop);
        float elapsed_time = 0.0f;
        cudaEventElapsedTime(&elapsed_time, start, stop);
        total_time += elapsed_time;
        
    }
    average_time = total_time / (test_count);
    printf("prop average time: %f ms\n",average_time);
    cudaDeviceSynchronize();
 
    //检验是否所有的结果都是3.0f。正确打印应该是0
    // TODO 如果去掉下面的cudaMemcpy会发生什么？
    cudaMemcpy(h_outputg, d_output, N*L*sizeof(ft), cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();
    cudaCheckErrors("cudaMemcpy failure");
    //for (int i = 0; i < N*L; i++) if (abs(h_outputg[i] - h_outputc[i])< h_outputc[i] * 0.01) {std::cout << "Mismatch at " << i << " was: " << h_outputg[i] << " should be: " << h_outputc[i] << std::endl; return average_time*1000;}
    std::cout << "Kernel execution time: " << dt/(float)USECPSEC << "s" << std::endl;

    //保存输出，用于后续比对
    std::ofstream ori_file(filename, std::ios::binary);
    ori_file.write(reinterpret_cast<char*>(h_outputg), sizeof(T) * N*L);
    ori_file.close();

    // 释放内存
    cudaDeviceSynchronize();
    cudaFree(d_output);
    cudaFree(d_input);
    cudaFree(d_matrix);

    // free(h_outputc);
    // free(h_outputg);
    free(h_input);
    free(h_matrix);

    return average_time;
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <output_filename>" << std::endl;
        return 1;
    }

    std::string base_filename = argv[1];
    // std::string base_filename = "/home/sw/jxie/work/cuda2maca/more_all_test/ppl_cukernel_clip_ndarray/output_ori/";

    std::vector<TestResult> results;

    // Test case 1
    std::string filename1 =  base_filename + "/test1.dat";
    float average_time1 = test_case<half>({2,128,256,512},{2,128,256,512},filename1);
    results.push_back({filename1, average_time1});

    // Test case 2
    std::string filename2 =  base_filename + "/test2.dat";
    float average_time2 = test_case<half>({2,256,256,1024},{2,256,256,1024},filename2);
    results.push_back({filename2, average_time2});

    // Test case 3
    std::string filename3 =  base_filename + "/test3.dat";
    float average_time3 = test_case<half>({3,128,256,2048},{3,128,256,2048},filename3);
    results.push_back({filename3, average_time3});

    // 每个测试样例保留一个输出结果用于判断优化一致性
    std::ofstream json_file(base_filename + "/test_results.json");
    json_file << "[\n";
    for(size_t i = 0; i < results.size(); ++i) {
        json_file << "  {\n";
        json_file << "    \"filename\": \"" << results[i].filename << "\",\n";
        json_file << "    \"average_time\": " << results[i].average_time << "\n";
        json_file << "  }";
        if(i < results.size() - 1) json_file << ",";
        json_file << "\n";
    }
    json_file << "]\n";
    json_file.close();

    return 0;
    
}